# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import cachetools
import contextlib
import json
import logging
import time

import coin.proto.coin_telemetry_pb2
from coin.base.timestamp import get_timestamp
from coin.proto.coin_strategy_pb2 import StrategyLog
from coin.proto.coin_telemetry_pb2 import (LatencyProto, LatencyMarkPoint, TelemetryProto)

LifeOfSignal = coin.proto.coin_telemetry_pb2


class Cache(cachetools.LRUCache):
  def __init__(
      self,
      maxsize,
      *,
      missing=None,
      getsizeof=None,
      evict_callback=None,
  ):
    super().__init__(maxsize, missing, getsizeof)
    self._evict_callback = evict_callback

  def popitem(self):
    key, value = super().popitem()
    if self._evict_callback:
      self._evict_callback(key, value)
    return key, value


class LatencyRecorder(object):
  def __init__(self, name: str, writer, feed_latency_record_freq=5):
    self._name = name
    self._context_id = int(time.time() * 1000)
    self._writer = writer
    self._records = Cache(maxsize=1024, evict_callback=self._evict_callback)
    self._logger = logging.getLogger('LatencyRecorder')
    self._feed_latency_record_freq = feed_latency_record_freq * 10**9
    self._last_feed_latency_ts = None

  @property
  def name(self):
    return self._name

  def get_current_context_id(self):
    return self._context_id

  def write(self, latency, flush=False):
    telemetry = TelemetryProto(type=TelemetryProto.LATENCY, latency=latency)
    log_proto = StrategyLog(type=StrategyLog.TELEMETRY, telemetry=telemetry)
    self._writer.write(log_proto, flush=flush)

  def begin(self, tag=None):
    self._context_id += 1
    self._records[self._context_id] = LatencyProto(name=self._name,
                                                   context_id=self._context_id,
                                                   tag=tag)

  def end(self, context_id=None):
    context_id = context_id or self._context_id
    record = self._records[context_id]
    record.done = True
    if record.tag == LatencyProto.FEED_LATENCY:
      current_ts = get_timestamp()
      if self._last_feed_latency_ts is not None and \
         current_ts - self._last_feed_latency_ts < self._feed_latency_record_freq:
        del self._records[context_id]
      else:
        self._last_feed_latency_ts = current_ts

  # noinspection PyUnusedLocal
  def _evict_callback(self, context_id, latency):
    # self._logger.debug('%s, %s', context_id, latency)
    if latency.done:
      self.write(latency)

  def set_time(self,
               point_name,
               *,
               account_request=None,
               pending_feed_number=None,
               context_id=None,
               timestamp=None,
               tag=None):
    timestamp = timestamp or get_timestamp()
    context_id = context_id or self._context_id
    if context_id not in self._records:
      self._logger.error('set_time: context_id %s is removed!', context_id)
      return
    mark_point = LatencyMarkPoint(point_name=point_name,
                                  account_request=account_request,
                                  pending_feed_number=pending_feed_number,
                                  context_id=context_id,
                                  timestamp=timestamp)
    if tag is not None:
      mark_point.tag = tag
    self._records[context_id].mark_points.extend([mark_point])
    self._records[context_id].last_update_timestamp = timestamp

  def set_extra_info(self, info: dict, *, context_id=None):
    context_id = context_id or self._context_id
    if context_id not in self._records:
      self._logger.error('set_extra_info: context_id %s is removed!', context_id)
      return
    self._records[context_id].extra_info.append(json.dumps(info))


@contextlib.contextmanager
def use_latency_recorder(latency_recorder, before_point_name, after_point_name, tag=None):
  if latency_recorder is not None:
    latency_recorder.set_time(before_point_name, tag=tag)
    yield
    latency_recorder.set_time(after_point_name, tag=tag)
  else:
    yield
